# -*- coding: utf-8 -*-
import glob
from matplotlib.colors import LogNorm
import xml.etree.ElementTree as ET
import matplotlib.pyplot as plt
import numpy as np
import sys
# plt.rcParams['font.sans-serif'] = ['YaHei Consolas Hybrid']  # 解决matplotlib无法输出中文，后面有教程
# plt.rcParams['axes.unicode_minus'] = False

def load_dataset(path):
    data = []
    plane = []
    sv = []
    ls = []
    kuang = []
    x = [134, 21, 24, 11, 16, 6, 44, 20, 10]
    y = [130, 18, 25, 6, 22, 11, 45, 11, 19]  # kmeans聚类结果
    for xml_file in glob.glob("{}/*xml".format(path)):
        tree = ET.parse(xml_file)
        filename = tree.findtext("filename")
        height = int(tree.findtext("./size/height"))
        width = int(tree.findtext("./size/width"))

        for obj in tree.iter("object"):
            name = obj.findtext("name")
            xmin = float(obj.findtext("bndbox/xmin"))
            ymin = float(obj.findtext("bndbox/ymin"))
            xmax = float(obj.findtext("bndbox/xmax"))
            ymax = float(obj.findtext("bndbox/ymax"))
            kuang.append([xmax - xmin, ymax - ymin])
            if xmax == xmin or ymax == ymin:
                print(xml_file)  # 上述为所有目标框一起可视化出来，下面是按类别可视化
            if name == "Railway Left":
                xmin = float(obj.findtext("bndbox/xmin"))
                ymin = float(obj.findtext("bndbox/ymin"))
                xmax = float(obj.findtext("bndbox/xmax"))
                ymax = float(obj.findtext("bndbox/ymax"))
                plane.append([xmax - xmin, ymax - ymin])
                if xmax == xmin or ymax == ymin:
                    print(xml_file)
            if name == "Railway Straight":
                xmin = float(obj.findtext("bndbox/xmin"))
                ymin = float(obj.findtext("bndbox/ymin"))
                xmax = float(obj.findtext("bndbox/xmax"))
                ymax = float(obj.findtext("bndbox/ymax"))
                sv.append([xmax - xmin, ymax - ymin])
                if xmax == xmin or ymax == ymin:
                    print(xml_file)
            if name == "Railway Right":
                xmin = float(obj.findtext("bndbox/xmin"))
                ymin = float(obj.findtext("bndbox/ymin"))
                xmax = float(obj.findtext("bndbox/xmax"))
                ymax = float(obj.findtext("bndbox/ymax"))
                ls.append([xmax - xmin, ymax - ymin])
                if xmax == xmin or ymax == ymin:
                    print(xml_file)

    pl = np.array(plane)
    sv1 = np.array(sv)
    ls1 = np.array(ls)
    kuang1 = np.array(kuang)

    plt.figure(1)
    plt.xlabel(u'X(像素)')
    plt.ylabel(u'Y(像素)')
    plt.title(u'plane')
    plt.hist2d(pl[:, 0], pl[:, 1], bins=400, norm=LogNorm())  # 画散点密度图
    plt.colorbar()
    # plt.scatter(pl[:,0],pl[:,1],s=2, c='k',marker = ".",alpha=0.4)
    plt.scatter(x, y, s=2, marker='.', c='r', alpha=0.4)
    plt.savefig(r'pl.png', dpi=300)
    plt.close(1)

    plt.figure(2)
    plt.xlabel(u'X(像素)')
    plt.ylabel(u'Y(像素)')
    plt.title(u'small-vehicle')
    plt.hist2d(sv1[:, 0], sv1[:, 1], bins=40, norm=LogNorm())
    plt.colorbar()
    # plt.scatter(sv1[:,0],sv1[:,1],s=2 ,c='k',marker = ".",alpha=0.8)
    plt.scatter(x, y, s=20, marker='.', c='r', alpha=0.4)
    plt.savefig(r'sv.png', dpi=300)
    plt.close(2)

    plt.figure(3)
    plt.xlabel(u'X(像素)')
    plt.ylabel(u'Y(像素)')
    plt.title(u'large-vehicle')
    plt.hist2d(ls1[:, 0], ls1[:, 1], bins=40, norm=LogNorm())
    plt.colorbar()
    # plt.scatter(ls1[:,0],ls1[:,1], s=2,c='k',marker = ".",alpha=0.4)
    plt.scatter(x, y, s=20, marker='.', c='r', alpha=0.4)
    plt.savefig(r'lv.png', dpi=300)
    plt.close(3)

    plt.figure(4)
    plt.xlabel(u'X(像素)')
    plt.ylabel(u'Y(像素)')
    plt.title(u'所有目标框')
    plt.hist2d(kuang1[:, 0], kuang1[:, 1], bins=400, norm=LogNorm())
    plt.colorbar()
    # plt.scatter(ls1[:,0],ls1[:,1], s=2,c='k',marker = ".",alpha=0.4)
    plt.scatter(x, y, s=2, marker='.', c='r', alpha=0.4)
    plt.savefig(r'all.png', dpi=300)
    plt.close(4)


if __name__ == '__main__':
    ANNOTATIONS_PATH = "E:\\Bullet\\VOCdevkit\\VOC2007\\Annotations"  # xml样本标签数据路径
    data = load_dataset(ANNOTATIONS_PATH)